#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库初始化和管理脚本
统一管理企业舆情分析系统的所有数据库表和索引
"""

import sqlite3
import json
import os
from datetime import datetime, date
from typing import Dict, List, Optional, Any


class DatabaseInitializer:
    """数据库初始化管理器"""
    
    def __init__(self, db_path: str = "enterprise_analysis.db"):
        self.db_path = db_path
        self.version = "1.0.0"
    
    def init_all_tables(self, recreate: bool = False):
        """初始化所有数据库表"""
        if recreate and os.path.exists(self.db_path):
            backup_path = f"{self.db_path}.backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            print(f"📁 备份现有数据库到: {backup_path}")
            os.rename(self.db_path, backup_path)
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        print("🚀 开始初始化数据库表...")
        
        try:
            # 1. 企业信息表
            self._create_enterprises_table(cursor)
            
            # 2. 分析会话表
            self._create_analysis_sessions_table(cursor)
            
            # 3. 问题记录表
            self._create_questions_table(cursor)
            
            # 4. AI响应表
            self._create_ai_responses_table(cursor)
            
            # 5. 威胁分析结果表
            self._create_threat_analysis_table(cursor)
            
            # 6. 威胁分析会话表
            self._create_threat_analysis_sessions_table(cursor)
            
            # 7. 威胁统计表
            self._create_threat_statistics_table(cursor)
            
            # 8. 系统配置表
            self._create_system_config_table(cursor)
            
            # 9. 审计日志表
            self._create_audit_log_table(cursor)
            
            # 创建索引
            self._create_indexes(cursor)
            
            # 创建视图
            self._create_views(cursor)
            
            # 插入初始配置
            self._insert_initial_config(cursor)
            
            conn.commit()
            print("✅ 数据库初始化完成！")
            
            # 显示表信息
            self.show_database_info()
            
        except Exception as e:
            conn.rollback()
            print(f"❌ 数据库初始化失败: {e}")
            raise
        finally:
            conn.close()
    
    def _create_enterprises_table(self, cursor):
        """创建企业信息表"""
        print("📊 创建企业信息表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS enterprises (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT UNIQUE NOT NULL,
                industry TEXT,
                description TEXT,
                competitors TEXT,  -- JSON数组
                keywords TEXT,     -- JSON数组
                founded_year INTEGER,
                headquarters TEXT,
                market_cap TEXT,
                business_scope TEXT,
                risk_level TEXT CHECK(risk_level IN ('低', '中', '高', 'low', 'medium', 'high')),
                metadata TEXT,     -- JSON对象
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
    
    def _create_analysis_sessions_table(self, cursor):
        """创建分析会话表"""
        print("📊 创建分析会话表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS analysis_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT UNIQUE NOT NULL,
                enterprise_name TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                status TEXT DEFAULT 'running' CHECK(status IN ('running', 'completed', 'error', 'cancelled')),
                total_questions INTEGER DEFAULT 0,
                successful_responses INTEGER DEFAULT 0,
                metadata TEXT,     -- JSON对象
                FOREIGN KEY (enterprise_name) REFERENCES enterprises (name)
            )
        ''')
    
    def _create_questions_table(self, cursor):
        """创建问题记录表"""
        print("📊 创建问题记录表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS questions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                question_text TEXT NOT NULL,
                question_type TEXT NOT NULL,
                is_inspirational BOOLEAN DEFAULT FALSE,
                generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                metadata TEXT,     -- JSON对象
                FOREIGN KEY (session_id) REFERENCES analysis_sessions (session_id)
            )
        ''')
    
    def _create_ai_responses_table(self, cursor):
        """创建AI响应表"""
        print("📊 创建AI响应表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS ai_responses (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                question_id INTEGER NOT NULL,
                ai_provider TEXT NOT NULL,
                response_text TEXT,
                success BOOLEAN DEFAULT TRUE,
                response_time REAL,
                error_message TEXT,
                usage_info TEXT,   -- JSON对象，存储token使用情况等
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (question_id) REFERENCES questions (id),
                FOREIGN KEY (session_id) REFERENCES analysis_sessions (session_id)
            )
        ''')
    
    def _create_threat_analysis_table(self, cursor):
        """创建威胁分析结果表"""
        print("📊 创建威胁分析结果表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS threat_analysis (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                question_id INTEGER NOT NULL,
                question_text TEXT NOT NULL,
                ai_provider TEXT NOT NULL,
                ai_response TEXT NOT NULL,
                threat_detected BOOLEAN NOT NULL,
                threat_types TEXT,  -- JSON数组存储威胁类型
                threat_level TEXT NOT NULL CHECK(threat_level IN ('low', 'medium', 'high', 'critical')),
                threat_score REAL NOT NULL CHECK(threat_score >= 0 AND threat_score <= 10),
                threat_description TEXT,
                recommended_actions TEXT,  -- JSON数组存储建议行动
                confidence_score REAL NOT NULL CHECK(confidence_score >= 0 AND confidence_score <= 1),
                analysis_details TEXT,  -- JSON存储详细分析信息
                analyzed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (question_id) REFERENCES questions (id),
                FOREIGN KEY (session_id) REFERENCES analysis_sessions (session_id)
            )
        ''')
    
    def _create_threat_analysis_sessions_table(self, cursor):
        """创建威胁分析会话表"""
        print("📊 创建威胁分析会话表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS threat_analysis_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                analysis_session_id TEXT UNIQUE NOT NULL,
                original_session_id TEXT NOT NULL,
                enterprise_name TEXT NOT NULL,
                total_responses INTEGER DEFAULT 0,
                threats_detected INTEGER DEFAULT 0,
                avg_threat_score REAL DEFAULT 0,
                highest_threat_level TEXT CHECK(highest_threat_level IN ('low', 'medium', 'high', 'critical')),
                analysis_started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                analysis_completed_at TIMESTAMP,
                status TEXT DEFAULT 'running' CHECK(status IN ('running', 'completed', 'error', 'cancelled')),
                metadata TEXT,     -- JSON对象
                FOREIGN KEY (original_session_id) REFERENCES analysis_sessions (session_id),
                FOREIGN KEY (enterprise_name) REFERENCES enterprises (name)
            )
        ''')
    
    def _create_threat_statistics_table(self, cursor):
        """创建威胁统计表"""
        print("📊 创建威胁统计表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS threat_statistics (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                enterprise_name TEXT NOT NULL,
                analysis_date DATE NOT NULL,
                total_threats INTEGER DEFAULT 0,
                critical_threats INTEGER DEFAULT 0,
                high_threats INTEGER DEFAULT 0,
                medium_threats INTEGER DEFAULT 0,
                low_threats INTEGER DEFAULT 0,
                avg_threat_score REAL DEFAULT 0,
                most_common_threat_type TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(enterprise_name, analysis_date),
                FOREIGN KEY (enterprise_name) REFERENCES enterprises (name)
            )
        ''')
    
    def _create_system_config_table(self, cursor):
        """创建系统配置表"""
        print("📊 创建系统配置表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS system_config (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                config_key TEXT UNIQUE NOT NULL,
                config_value TEXT,
                config_type TEXT DEFAULT 'string' CHECK(config_type IN ('string', 'json', 'number', 'boolean')),
                description TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
    
    def _create_audit_log_table(self, cursor):
        """创建审计日志表"""
        print("📊 创建审计日志表...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS audit_log (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                action TEXT NOT NULL,
                table_name TEXT,
                record_id TEXT,
                user_info TEXT,
                changes TEXT,      -- JSON对象记录变更详情
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                ip_address TEXT,
                user_agent TEXT
            )
        ''')
    
    def _create_indexes(self, cursor):
        """创建数据库索引"""
        print("🔗 创建数据库索引...")
        
        indexes = [
            # 企业信息表索引
            "CREATE INDEX IF NOT EXISTS idx_enterprises_name ON enterprises (name)",
            "CREATE INDEX IF NOT EXISTS idx_enterprises_industry ON enterprises (industry)",
            "CREATE INDEX IF NOT EXISTS idx_enterprises_risk_level ON enterprises (risk_level)",
            
            # 分析会话表索引
            "CREATE INDEX IF NOT EXISTS idx_analysis_sessions_session_id ON analysis_sessions (session_id)",
            "CREATE INDEX IF NOT EXISTS idx_analysis_sessions_enterprise ON analysis_sessions (enterprise_name)",
            "CREATE INDEX IF NOT EXISTS idx_analysis_sessions_status ON analysis_sessions (status)",
            "CREATE INDEX IF NOT EXISTS idx_analysis_sessions_created_at ON analysis_sessions (created_at)",
            
            # 问题表索引
            "CREATE INDEX IF NOT EXISTS idx_questions_session_id ON questions (session_id)",
            "CREATE INDEX IF NOT EXISTS idx_questions_type ON questions (question_type)",
            "CREATE INDEX IF NOT EXISTS idx_questions_inspirational ON questions (is_inspirational)",
            
            # AI响应表索引
            "CREATE INDEX IF NOT EXISTS idx_ai_responses_session_id ON ai_responses (session_id)",
            "CREATE INDEX IF NOT EXISTS idx_ai_responses_question_id ON ai_responses (question_id)",
            "CREATE INDEX IF NOT EXISTS idx_ai_responses_provider ON ai_responses (ai_provider)",
            "CREATE INDEX IF NOT EXISTS idx_ai_responses_success ON ai_responses (success)",
            
            # 威胁分析表索引
            "CREATE INDEX IF NOT EXISTS idx_threat_analysis_session_id ON threat_analysis (session_id)",
            "CREATE INDEX IF NOT EXISTS idx_threat_analysis_question_id ON threat_analysis (question_id)",
            "CREATE INDEX IF NOT EXISTS idx_threat_analysis_detected ON threat_analysis (threat_detected)",
            "CREATE INDEX IF NOT EXISTS idx_threat_analysis_level ON threat_analysis (threat_level)",
            "CREATE INDEX IF NOT EXISTS idx_threat_analysis_score ON threat_analysis (threat_score)",
            
            # 威胁分析会话表索引
            "CREATE INDEX IF NOT EXISTS idx_threat_sessions_analysis_id ON threat_analysis_sessions (analysis_session_id)",
            "CREATE INDEX IF NOT EXISTS idx_threat_sessions_original_id ON threat_analysis_sessions (original_session_id)",
            "CREATE INDEX IF NOT EXISTS idx_threat_sessions_enterprise ON threat_analysis_sessions (enterprise_name)",
            
            # 威胁统计表索引
            "CREATE INDEX IF NOT EXISTS idx_threat_stats_enterprise_date ON threat_statistics (enterprise_name, analysis_date)",
            
            # 系统配置表索引
            "CREATE INDEX IF NOT EXISTS idx_system_config_key ON system_config (config_key)",
            
            # 审计日志表索引
            "CREATE INDEX IF NOT EXISTS idx_audit_log_action ON audit_log (action)",
            "CREATE INDEX IF NOT EXISTS idx_audit_log_table ON audit_log (table_name)",
            "CREATE INDEX IF NOT EXISTS idx_audit_log_timestamp ON audit_log (timestamp)"
        ]
        
        for index_sql in indexes:
            cursor.execute(index_sql)
    
    def _create_views(self, cursor):
        """创建数据库视图"""
        print("👁️ 创建数据库视图...")
        
        # 企业分析概览视图
        cursor.execute('''
            CREATE VIEW IF NOT EXISTS enterprise_analysis_overview AS
            SELECT 
                e.name as enterprise_name,
                e.industry,
                e.risk_level,
                COUNT(DISTINCT s.session_id) as total_sessions,
                COUNT(DISTINCT q.id) as total_questions,
                COUNT(DISTINCT r.id) as total_responses,
                AVG(r.response_time) as avg_response_time,
                MAX(s.created_at) as last_analysis_date
            FROM enterprises e
            LEFT JOIN analysis_sessions s ON e.name = s.enterprise_name
            LEFT JOIN questions q ON s.session_id = q.session_id
            LEFT JOIN ai_responses r ON q.id = r.question_id AND r.success = 1
            GROUP BY e.name, e.industry, e.risk_level
        ''')
        
        # 威胁分析概览视图
        cursor.execute('''
            CREATE VIEW IF NOT EXISTS threat_analysis_overview AS
            SELECT 
                ta.session_id,
                s.enterprise_name,
                COUNT(*) as total_responses_analyzed,
                SUM(CASE WHEN ta.threat_detected THEN 1 ELSE 0 END) as threats_detected,
                AVG(ta.threat_score) as avg_threat_score,
                MAX(ta.threat_level) as highest_threat_level,
                COUNT(CASE WHEN ta.threat_level = 'critical' THEN 1 END) as critical_threats,
                COUNT(CASE WHEN ta.threat_level = 'high' THEN 1 END) as high_threats,
                COUNT(CASE WHEN ta.threat_level = 'medium' THEN 1 END) as medium_threats,
                COUNT(CASE WHEN ta.threat_level = 'low' THEN 1 END) as low_threats
            FROM threat_analysis ta
            JOIN analysis_sessions s ON ta.session_id = s.session_id
            GROUP BY ta.session_id, s.enterprise_name
        ''')
        
        # AI提供商性能视图
        cursor.execute('''
            CREATE VIEW IF NOT EXISTS ai_provider_performance AS
            SELECT 
                ai_provider,
                COUNT(*) as total_requests,
                SUM(CASE WHEN success THEN 1 ELSE 0 END) as successful_requests,
                ROUND(AVG(CASE WHEN success THEN 1.0 ELSE 0.0 END) * 100, 2) as success_rate,
                ROUND(AVG(response_time), 3) as avg_response_time,
                MIN(response_time) as min_response_time,
                MAX(response_time) as max_response_time
            FROM ai_responses
            GROUP BY ai_provider
        ''')
    
    def _insert_initial_config(self, cursor):
        """插入初始系统配置"""
        print("⚙️ 插入初始系统配置...")
        
        configs = [
            ('db_version', self.version, 'string', '数据库版本'),
            ('init_date', datetime.now().isoformat(), 'string', '数据库初始化日期'),
            ('max_concurrent_requests', '10', 'number', '最大并发请求数'),
            ('default_timeout', '30', 'number', '默认超时时间（秒）'),
            ('enable_audit_log', 'true', 'boolean', '是否启用审计日志'),
            ('threat_score_threshold', '7.0', 'number', '威胁评分阈值'),
            ('supported_ai_providers', '["baidu", "aliyun", "zhipu", "kimi", "deepseek"]', 'json', '支持的AI提供商'),
            ('question_types', '["negative_news", "positive_news", "financial_info", "market_analysis", "product_info", "leadership", "competitive_analysis", "technology", "reputation", "investment", "regulation", "esg"]', 'json', '支持的问题类型'),
            ('threat_types', '["negative_sentiment", "data_leak", "reputation_damage", "financial_risk", "legal_risk", "operational_risk", "security_breach", "compliance_violation", "competitive_disadvantage"]', 'json', '威胁类型列表')
        ]
        
        for config_key, config_value, config_type, description in configs:
            cursor.execute('''
                INSERT OR IGNORE INTO system_config (config_key, config_value, config_type, description)
                VALUES (?, ?, ?, ?)
            ''', (config_key, config_value, config_type, description))
    
    def show_database_info(self):
        """显示数据库信息"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        print(f"\n{'='*60}")
        print(f"📊 数据库信息 ({self.db_path})")
        print(f"{'='*60}")
        
        # 获取所有表
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
        tables = cursor.fetchall()
        
        print(f"📋 数据表 ({len(tables)} 个):")
        for table in tables:
            table_name = table[0]
            cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
            count = cursor.fetchone()[0]
            print(f"   • {table_name:25} | {count:6} 条记录")
        
        # 获取所有视图
        cursor.execute("SELECT name FROM sqlite_master WHERE type='view' ORDER BY name")
        views = cursor.fetchall()
        
        if views:
            print(f"\n👁️ 数据视图 ({len(views)} 个):")
            for view in views:
                print(f"   • {view[0]}")
        
        # 获取所有索引
        cursor.execute("SELECT name FROM sqlite_master WHERE type='index' AND name NOT LIKE 'sqlite_%' ORDER BY name")
        indexes = cursor.fetchall()
        
        if indexes:
            print(f"\n🔗 索引 ({len(indexes)} 个):")
            for index in indexes:
                print(f"   • {index[0]}")
        
        # 数据库文件大小
        file_size = os.path.getsize(self.db_path)
        if file_size < 1024:
            size_str = f"{file_size} B"
        elif file_size < 1024 * 1024:
            size_str = f"{file_size / 1024:.1f} KB"
        else:
            size_str = f"{file_size / (1024 * 1024):.1f} MB"
        
        print(f"\n💾 数据库大小: {size_str}")
        
        conn.close()
    
    def validate_database(self) -> bool:
        """验证数据库完整性"""
        print("🔍 验证数据库完整性...")
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            # 检查PRAGMA完整性
            cursor.execute("PRAGMA integrity_check")
            result = cursor.fetchone()
            
            if result[0] != "ok":
                print(f"❌ 数据库完整性检查失败: {result[0]}")
                return False
            
            # 检查外键约束
            cursor.execute("PRAGMA foreign_key_check")
            fk_errors = cursor.fetchall()
            
            if fk_errors:
                print(f"❌ 外键约束错误: {len(fk_errors)} 个")
                for error in fk_errors:
                    print(f"   {error}")
                return False
            
            print("✅ 数据库完整性验证通过")
            return True
            
        except Exception as e:
            print(f"❌ 数据库验证异常: {e}")
            return False
        finally:
            conn.close()
    
    def backup_database(self, backup_path: str = None) -> str:
        """备份数据库"""
        if backup_path is None:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            backup_path = f"{self.db_path}.backup_{timestamp}"
        
        import shutil
        shutil.copy2(self.db_path, backup_path)
        print(f"💾 数据库已备份到: {backup_path}")
        return backup_path
    
    def clean_old_data(self, days: int = 30):
        """清理旧数据"""
        print(f"🧹 清理 {days} 天前的旧数据...")
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            cutoff_date = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
            cutoff_timestamp = cutoff_date.timestamp() - (days * 24 * 3600)
            cutoff_str = datetime.fromtimestamp(cutoff_timestamp).isoformat()
            
            # 清理审计日志
            cursor.execute("DELETE FROM audit_log WHERE timestamp < ?", (cutoff_str,))
            audit_deleted = cursor.rowcount
            
            # 清理已完成的旧分析会话（保留威胁分析结果）
            cursor.execute("""
                DELETE FROM analysis_sessions 
                WHERE status = 'completed' AND created_at < ?
                AND session_id NOT IN (
                    SELECT original_session_id FROM threat_analysis_sessions
                )
            """, (cutoff_str,))
            sessions_deleted = cursor.rowcount
            
            conn.commit()
            
            print(f"✅ 清理完成:")
            print(f"   • 审计日志: {audit_deleted} 条")
            print(f"   • 分析会话: {sessions_deleted} 个")
            
        except Exception as e:
            conn.rollback()
            print(f"❌ 数据清理失败: {e}")
        finally:
            conn.close()


if __name__ == "__main__":
    print("🎯 企业舆情分析系统数据库初始化工具")
    print("=" * 60)
    
    db_init = DatabaseInitializer()
    
    if not os.path.exists(db_init.db_path):
        print("📁 数据库文件不存在，正在创建...")
        db_init.init_all_tables()
    else:
        print("📁 数据库文件已存在")
        db_init.show_database_info()
        
        choice = input("\n是否重新初始化数据库？(y/N): ").strip().lower()
        if choice in ['y', 'yes']:
            db_init.init_all_tables(recreate=True)
        elif input("是否验证数据库完整性？(Y/n): ").strip().lower() not in ['n', 'no']:
            db_init.validate_database()
    
    print("\n✅ 数据库初始化脚本执行完成！") 